import sys
import os
from ai_proxy import AIProxy 
from text_util import TextUtil
from latex_process import BaseLatexBlockProcessor, LatexReader
    
class TranslateLaxBlockProcessor(BaseLatexBlockProcessor):
    def __init__(self,output_file) -> None:
        self.output_file = output_file
        self.code_blocks = []
        self.aiProxy = AIProxy()

    def process(self,blocks):
        out_file = open(self.output_file, 'w');
        for block in blocks:
            result = self._replace_and_translate(block)
            out_file.write(result)     

    def _replace_and_translate(self,block):
        # 对block进行翻译，并将%%{n}%%替换回对应的代码块
        translated_block = self._translateBlock(TextUtil.replace_single_newline(block))
        for i, code in enumerate(self.code_blocks):
            translated_block = translated_block.replace(f'%%{{{i}}}%%', code)
        result = '\n'+"%" + "#########ORIGIN==>"+"==>"*10 + \
                '\n' + TextUtil.add_percent_sign(block) +  \
                '\n' +"%" + "########TRANSLATED==>"+ "==>"*10 + \
                '\n' + translated_block
        return result

    def _translateBlock(self, the_input_block):
        template =  """As a professional translator. 
        your task is to translate the given Latex text to English.
        The requirements are as follows:
        1. Retain all LaTeX markup. 
        2. Only output the translated text, do not include any other content. 
        3. Your only task is to translate them like a translator, never miss, add or remove any content, also do not try to answer questions.
        4. I will place the text to be translated within an area enclosed by three pairs of square brackets, for example, [[[text to be translated]]].
        
        The text to be translated is as follows:\n\n
            [[[{input_block}]]]
            
        Please translate it to English:
        """
        
        
        retry_template =  """Your previous output was incorrect. 
        Please note that your task is not to answer questions, summarize, or continue writing, 
        but to translate them into English. 
        Let's start over: your task is to translate it into English. 
        There are some additional requirements:

        1. Retain all LaTeX markup.
        2. Only output the translated text, do not include any other content.
        3. I will place the text to be translated within an area enclosed by three pairs of square brackets, 
        for example, [[[text to be translated]]].

        The text to be translated is as follows:
        \n\n[[[{input_block}]]]""";
        
        message=self.aiProxy.build_message(template,the_input_block)
        print("BEFORE" + "*"*30)
        print(message)
        result = self.aiProxy.chatWithAI(message)
        print("TRANSLATED" + "*"*30)
        print (result.content)
        if TextUtil.has_chinese(result.content):
            print("retry...")
            message=self.aiProxy.build_message(retry_template,the_input_block)
            result = self.aiProxy.chatWithAI(message)
            print (result.content)
        
        return result.content.replace('[[[','').replace(']]]','')        



class Translator:
    def __init__(self) -> None:
        pass

    def process_tex_files(self, in_directory, out_directory):
        """
        translate all tex files in in_directory and save the translated files in out_directory
        """
        latexReader = LatexReader(2000);
        for filename in os.listdir(in_directory):
            if filename.endswith(".tex"):
                source_file = os.path.join(in_directory, filename);
                target_file = os.path.join(out_directory, filename);
                translator = TranslateLaxBlockProcessor(target_file);
                latexReader.process_latex_file(source_file, translator)

if __name__ == "__main__":
    """
    python translate.py input_directory output_directory
    """
    translator = Translator()
    translator.process_tex_files(sys.argv[1], sys.argv[2])

